# ---------------------------------- #
# Virtual Machine Snapshot Functions #
# ---------------------------------- #

# Create virtual machine snapshots
function SnapshotCreate-VirtualMachines([System.Xml.XmlElement]$subscription)
{
	if (!(Should-Run($subscription.virtualMachines))) {return}
    Load-StorageClient
	foreach ($virtualMachine in $subscription.virtualMachines.virtualMachine)
	{
	    if (!(Should-Run($virtualMachine))) {continue}
		SnapshotCreate-VirtualMachineInstance $virtualMachine
    }
}

# List virtual machine snapshots
function SnapshotList-VirtualMachines([System.Xml.XmlElement]$subscription)
{
	if (!(Should-Run($subscription.virtualMachines))) {return}
    Load-StorageClient
	foreach ($virtualMachine in $subscription.virtualMachines.virtualMachine)
	{
	    if (!(Should-Run($virtualMachine))) {continue}
		SnapshotList-VirtualMachineInstance $virtualMachine
    }
}

# Remove virtual machine snapshots
function SnapshotRemove-VirtualMachines([System.Xml.XmlElement]$subscription)
{
	if (!(Should-Run($subscription.virtualMachines))) {return}
    Load-StorageClient
	foreach ($virtualMachine in $subscription.virtualMachines.virtualMachine)
	{
	    if (!(Should-Run($virtualMachine))) {continue}
		SnapshotRemove-VirtualMachineInstance $virtualMachine
    }
}

# Restore virtual machine snapshots
function SnapshotRestore-VirtualMachines([System.Xml.XmlElement]$subscription)
{
	if (!(Should-Run($subscription.virtualMachines))) {return}
    Load-StorageClient
	foreach ($virtualMachine in $subscription.virtualMachines.virtualMachine)
	{
	    if (!(Should-Run($virtualMachine))) {continue}
		SnapshotRestore-VirtualMachineInstance $virtualMachine
    }
}

# Create virtual machine instance snapshot
function SnapshotCreate-VirtualMachineInstance($virtualMachine)
{
    $vm = Get-AzureVM -ServiceName $virtualMachine.service -Name $virtualMachine.name -WarningAction Ignore
	if ($vm -ne $null)
	{
        $started = $vm.PowerState -eq "Started"
        if ($started)
        {
            Stop-VirtualMachineInstance $virtualMachine $true $true
        }
		Write-Information "Creating virtual machine '$($virtualMachine.name)' snapshot ... " -NoNewLine $true
        $options = New-Object Microsoft.WindowsAzure.StorageClient.BlobRequestOptions
	    foreach ($disk in (Get-AzureDisk | ? {$_.AttachedTo.HostedServiceName -eq $virtualMachine.service -and $_.AttachedTo.RoleName -eq $virtualMachine.name} | Sort-Object -Property Label))
	    {
            $client = Get-StorageClient $disk
		    $blob = $client.GetBlobReference($disk.MediaLink.LocalPath.Substring(1))
		    $snapshot = $blob.CreateSnapshot($options)
	    }
		Write-Text "done."
        if ($started)
        {
            Start-VirtualMachineInstance $virtualMachine
        }
	}
	else
	{
		Write-Information "Creating virtual machine '$($virtualMachine.name)' snapshot ... skipped."
	}
}

# List virtual machine instance snapshots
function SnapshotList-VirtualMachineInstance($virtualMachine)
{
	if ((Get-AzureVM -ServiceName $virtualMachine.service -Name $virtualMachine.name -WarningAction Ignore) -ne $null)
	{
		Write-Information "Listing virtual machine '$($virtualMachine.name)' snapshot ... " -NoNewLine $true
        $snapshots = @()
	    foreach ($disk in (Get-AzureDisk | ? {$_.AttachedTo.HostedServiceName -eq $virtualMachine.service -and $_.AttachedTo.RoleName -eq $virtualMachine.name} | Sort-Object -Property Label))
	    {
            $snapshots += Get-DiskSnapshots $disk
        }
		Write-Text "done."
        $snapshots | Sort-Object -Property Name, SnapshotTime | Format-Table -Property Name, SnapshotTime -AutoSize
	}
	else
	{
		Write-Information "Listing virtual machine '$($virtualMachine.name)' snapshot ... skipped."
	}
}

# Remove virtual machine instance snapshot
function SnapshotRemove-VirtualMachineInstance($virtualMachine)
{
	if ((Get-AzureVM -ServiceName $virtualMachine.service -Name $virtualMachine.name -WarningAction Ignore) -ne $null)
	{
		Write-Information "Removing virtual machine '$($virtualMachine.name)' snapshot ... " -NoNewLine $true
	    foreach ($disk in (Get-AzureDisk | ? {$_.AttachedTo.HostedServiceName -eq $virtualMachine.service -and $_.AttachedTo.RoleName -eq $virtualMachine.name} | Sort-Object -Property Label))
	    {
		    $snapshot = Get-DiskSnapshots $disk | Sort-Object -Property SnapshotTime -Descending | Select-Object -First 1
		    if ($snapshot -ne $null)
		    {
			    $snapshot.Delete()
		    }
	    }
		Write-Text "done."
	}
	else
	{
		Write-Information "Removing virtual machine '$($virtualMachine.name)' snapshot ... skipped."
	}
}

# Restore virtual machine instance snapshot
function SnapshotRestore-VirtualMachineInstance($virtualMachine)
{
	if ((Get-AzureVM -ServiceName $virtualMachine.service -Name $virtualMachine.name -WarningAction Ignore) -ne $null)
	{
        $started = $vm.PowerState -eq "Started"
        Stop-VirtualMachineInstance $virtualMachine $true $true
		Write-Information "Restoring virtual machine '$($virtualMachine.name)' snapshot ... " -NoNewLine $true
	    $file = "$env:dp0\VM_$($virtualMachine.name).xml"
	    $operation = Export-AzureVM -ServiceName $virtualMachine.service -Name $virtualMachine.name -Path $file
	    if (!$? -or $operation.OperationStatus -ne "Succeeded") { throw "Virtual machine '$($virtualMachine.name)' operation failed." }
        $disks = Get-AzureDisk | ? {$_.AttachedTo.HostedServiceName -eq $virtualMachine.service -and $_.AttachedTo.RoleName -eq $virtualMachine.name} | Sort-Object -Property Label
        $fileDisks = Save-DisksToFile $disks
	    Remove-VirtualMachineInstanceConfiguration $virtualMachine $true
        foreach ($disk in $disks)
        {
            $snapshot = Get-DiskSnapshots $disk | Sort-Object -Property SnapshotTime -Descending | Select-Object -First 1
            if ($snapshot -ne $null)
            {
                $client = Get-StorageClient $disk
                $blob = $client.GetBlobReference($disk.MediaLink.LocalPath.Substring(1))
                $result = $blob.CopyFromBlob($snapshot)
            }
            if ($disk.OS -ne $null)
            {
                $operation = Add-AzureDisk -DiskName $disk.DiskName -MediaLocation $disk.MediaLink -OS $disk.OS
            }
            else
            {
                $operation = Add-AzureDisk -DiskName $disk.DiskName -MediaLocation $disk.MediaLink
            }
            if (!$? -or $operation.OperationStatus -ne "Succeeded") { throw "Virtual machine '$($virtualMachine.name)' operation failed." }
        }
	    $operation = Import-AzureVM -Path $file | New-AzureVM -ServiceName $virtualMachine.service -VNetName $virtualMachine.network.name -WarningAction Ignore
	    if (!$? -or $operation.OperationStatus -ne "Succeeded") { throw "Virtual machine '$($virtualMachine.name)' operation failed." }
        Remove-Item $fileDisks -Force
	    Remove-Item $file -Force
		Write-Text "done."
	}
	else
	{
		Write-Information "Restoring virtual machine '$($virtualMachine.name)' snapshot ... skipped."
	}
}

# Load Storage Client library
function Load-StorageClient
{
	Add-Type -Path "C:\Program Files\Microsoft SDKs\Azure\.NET SDK\v2.5\bin\Microsoft.WindowsAzure.StorageClient.dll"
}

# Return Storage Client instance
function Get-StorageClient($disk)
{
	$account = Get-DiskAccountName $disk
	$uri = "$($disk.MediaLink.Scheme)://$account.blob.core.windows.net/"
	$credentials = New-Object Microsoft.WindowsAzure.StorageCredentialsAccountAndKey($account, (Get-AzureStorageKey -StorageAccountName $account).Primary)
	return New-Object Microsoft.WindowsAzure.StorageClient.CloudBlobClient($uri, $credentials)
}

# Return disk account name
function Get-DiskAccountName($disk)
{
    return $disk.MediaLink.Host.Substring(0, $disk.MediaLink.Host.IndexOf("."))
}

# Return disk snapshots
function Get-DiskSnapshots($disk)
{
	$options = New-Object Microsoft.WindowsAzure.StorageClient.BlobRequestOptions
	$options.UseFlatBlobListing = $true
	$options.BlobListingDetails = [Microsoft.WindowsAzure.StorageClient.BlobListingDetails]::Snapshots
    $client = Get-StorageClient $disk
    $account = Get-DiskAccountName $disk
    $container = $client.GetContainerReference("$($disk.MediaLink.Scheme)://$account.blob.core.windows.net/$($disk.MediaLink.Segments[1])")
    return $container.ListBlobs($options) | ? {$_.SnapshotTime -ne $null -and $_.Name -eq $disk.MediaLink.Segments[2]}
}

# Save disks information to file (in case of processing exception, otherwise the Virtual Machine instance configuration will be lost)
function Save-DisksToFile($disks)
{
    $fileName = "$env:dp0\VM_$($virtualMachine.name)_Disks.xml"
    (ConvertTo-Xml -InputObject $disks).Save($fileName)
    return $fileName
}